# function definitions
from lxml import etree
import file_ops, xml

# custom import (C-like include)
# geckoloader_path and wit_path will be defined there
exec(open("bin_tool.py").read())

# function to return an attribute string with the
# {$__gameid}, {$__region}, {$__maker} string replacements done
def rtn_def_param_str(string, params_list):
  
  # return this!
  return file_ops.get_path_str(string.replace("{$__gameid}", params_list[0]).replace("{$__region}", params_list[1]).replace("{$__maker}", params_list[2]))

# function to check if the game image is valid
def check_game_image(game_image_path):
  result = exec_subprocess([wit_path, "verify", "--test", game_image_path])
  if (result.returncode != 0):
    return False
  return True
    
# function to check if the game image is valid for the riivolution XML
def check_game_xml_id(game_image_path, riiv_xml_id):
  
  # get the actual game id
  game_id = exec_subprocess([wit_path, "ID6", game_image_path]).stdout.replace("\n", "")
  
  # riiv_xml_id will be formatted as
  # ["XXX", ["X", "X", ...], "XX"]
  #   RMG     E    P   ...    01    (example)
  
  # see which parts of the game id are valid
  is_id_sec_valid = [False, False, False]
  
  # first 3 letters
  if (riiv_xml_id[0] == game_id[:-3] or riiv_xml_id[0] == ""):
    is_id_sec_valid[0] = True
  # 4th letter
  if (len(riiv_xml_id[1]) == 0):
    is_id_sec_valid[1] = True
  else:
    for reg in riiv_xml_id[1]:
      if (reg == game_id[3:-2]):
        is_id_sec_valid[1] = True
  # last 2 letters
  if (riiv_xml_id[2] == game_id[4:] or riiv_xml_id[2] == ""):
    is_id_sec_valid[2] = True
  
  # verify all parts 
  if (False in is_id_sec_valid):
    return False
  return game_id

# check_riiv_xml() function
# function to check if a riivolution XML is valid using a XSD file
# if the function returns False and the file actually seems right 
# then you have to use a service like the one in this page to
# manually check what is the issue
def check_riiv_xml(xml_path):  
  
  # check agaisnt MyRiivolution.xsd
  return xml.check_with_sch(xml_path, "MyRiivolution.xsd")

# get_patch_elem_root_path() function
# function to get the root path of the files inside a patch element
def get_patch_elem_root_path(root_elem, patch_elem):

  # get the root path of the whole XML
  root_path = ""
  if ("root" in root_elem.attrib):
    root_path = root_path + root_elem.attrib["root"]
  # get the relative path of the patch element
  if ("root" in patch_elem.attrib):
    root_path = root_path + patch_elem.attrib["root"]
  return file_ops.get_path_str(root_path)

# check_riiv_patches() function
# function to check the information about the patches to be applied to the game
# this function will basically check if filepaths are valid both for the XML and the game
def check_riiv_patches(game_path, xml_path, mod_files_folder):  
    
  ##################################
  # check the game file and xml file
  if (check_game_image(game_path) == False):
    print("Invalid ISO/WBFS")
    return False
  if (check_riiv_xml(xml_path) == False):
    print("Invalid Riivolution XML")
    return False  
  
  # check if the game id matches the game id from the riivolution file
  riiv_game_id = ["", [], ""]
  xml = etree.parse(xml_path, etree.XMLParser(ns_clean = True, remove_comments = True)) # parse the xml file
  root = xml.getroot()
  
  #########################################
  # check the id tag in the riivolution XML
  id_elem = root.find("id")
  if (id_elem != None):
    # grab the game id tag
    if ("game" in id_elem.attrib):
      # game part
      riiv_game_id[0] = id_elem.attrib["game"]
      # region part
      if (len(riiv_game_id[0]) == 3): # check if the game id is 3 characters long
        if (len(id_elem) != 0):
          for reg in id_elem:
            riiv_game_id[1].append(reg.attrib["type"])
      else: # it is 4 characters long
        riiv_game_id[1].append(riiv_game_id[0][3:])
        riiv_game_id[0] = riiv_game_id[0][:-1]
      # developer part
      if ("developer" in id_elem.attrib):
        riiv_game_id[2] = id_elem.attrib["developer"]
  
  # check all the possible combinations of a full game id
  game_id = ["", "", ""]
  game_id[0] = check_game_xml_id(game_path, riiv_game_id)
  if (game_id == False):
    print("Riivolution XML is not compatible with the Wii game provided")
    return False
  game_id[1] = game_id[0][3:-2]
  game_id[2] = game_id[0][4:]
  game_id[0] = game_id[0][:-3]
  
  
  ######################################################################
  # parse the XML file and check the patches and their file replacements
  opt_patches_id = []
  for child in root:
    if (child.tag == "options"):
      for sec in child:
        for opt in sec:
          for choice in opt:
            for patch in choice:
              opt_patches_id.append(patch.attrib["id"]) # lol
  elems_patch_id = []
  for patch in root.findall("patch"):
    elems_patch_id.append(patch.attrib["id"])
  # check if all the options patches id exist in the patch tags down in the XML
  patch_elems = []
  for patch_id in opt_patches_id:
    if (patch_id not in elems_patch_id):
      print("Patch to apply from options does not exist in the XML")
      return False
    # at the same time store the patch elements that will actually get read by riivolution
    for patch in root.findall("patch"):
      if (patch.attrib["id"] == patch_id):
        patch_elems.append(patch)  
  
  ######################################################
  # check the patch tags file/folder/memory replacements
  # but before that, extract the game files list (yikes)
  game_files = exec_subprocess([wit_path, "files", game_path]).stdout.split("\n")
  # remove non-used elements
  cont_loop = True
  while (cont_loop == True):
    for i in range(0, len(game_files)):
      if (("DATA/files" not in game_files[i]) and ("./files" not in game_files[i])):
        game_files.pop(i)
        cont_loop = True
        break
      cont_loop = False
  # do another for loop to get rid of "DATA/files/" and "./files/" starting strings
  for i in range(0, len(game_files)):
    if ("DATA/files" in game_files[i]):
      game_files[i] = file_ops.get_path_str(game_files[i].replace("DATA/files", ""))
    else:
      game_files[i] = file_ops.get_path_str(game_files[i].replace("./files", ""))
  
  # ~ for i in range(0, len(game_files)):
    # ~ print(game_files[i])
  
  ##########################################
  # start the loop! and check the file paths
  for patch in patch_elems:
    # get the root path of the current pacthes
    patch_root_path = rtn_def_param_str(file_ops.get_base_path(mod_files_folder, True) + "/"
                                        + get_patch_elem_root_path(root, patch), game_id)
    # check the folder, file tag patches
    # Note: the offset, length and other attributes for file and folder
    #       elements are weird and I don't know how to support them
    for patch_type in patch:
      # file patch
      if (patch_type.tag == "file"):
        fpath = patch_root_path + "/" + rtn_def_param_str(patch_type.attrib["external"], game_id)
        if (file_ops.f_exists(fpath) == False):
          print("Line %d: File patch has invalid external path" % (patch_type.sourceline))
          print(fpath)
          return False
        else: # check if the file exists on disc
          if (patch_type.attrib["disc"] not in game_files):
            if ("create" not in patch_type.attrib):
              print("Line %d: File patch creates non-existant disc file without create attribute" % (patch_type.sourceline))
              return False
            else: # create attribute exists
              if (patch_type.attrib["create"] == "false"):
                print("Line %d: File patch creates non-existant disc file with create = False" % (patch_type.sourceline))
                return False
      
      # folder patch
      if (patch_type.tag == "folder"):
        fpath = patch_root_path + "/" + rtn_def_param_str(patch_type.attrib["external"], game_id)
        if (file_ops.f_exists(fpath) == False):
          print("Line %d: Folder patch has invalid external path" % (patch_type.sourceline))
          print(fpath)
          return False
        if ("disc" in patch_type.attrib):
          if (rtn_def_param_str(patch_type.attrib["disc"], game_id) not in game_files):
            # check if the create attribute is present
            if ("create" in patch_type.attrib):
              if (patch_type.attrib["create"] == "False"):
                print("Line %d: Folder patch creates non-existant disc folder with create = False" % (patch_type.sourceline))
                return False
            else: # cannot create a disc non-existant folder without a create attribute
              print("Line %d: Folder patch creates non-existant disc folder without create attribute" % (patch_type.sourceline))
              return False
      # memory patch (incomplete, full check happens when building the mod)
      if (patch_type.tag == "memory"):
        if ("valuefile" in patch_type.attrib):
          fpath = patch_root_path + "/" + rtn_def_param_str(patch_type.attrib["valuefile"], game_id)
          if (file_ops.f_exists(fpath) == False):
            print("Line %d: Memory patch has invalid paths")
            print(fpath)
            return False
        # restrict patches to be applied to the 0x80000000 memory area
        if ("0x8" not in patch_type.attrib["offset"]):
          print("Line %d: Memory patch tries patching to a range outside the 0x80000000 memory area" % (patch_type.sourceline))
          return False
        # check if a ram dump is needed for the patching process
        if ("original" in patch_type.attrib):
          if (file_ops.f_exists("ram_dumps/" + game_id[0] + game_id[1] + game_id[2] + ".bin") == False
              or file_ops.get_file_size("ram_dumps/" + game_id[0] + game_id[1] + game_id[2] + ".bin") != 0x1800000):
            print("Line %d: Memory element has original condition." % (patch_type.sourceline))
            print("No/Invalid ram dump of \"%s\" was found in ram_dumps/" % (game_id[0] + game_id[1] + game_id[2]))
            return False
  
  ###############################
  # hardcore memory patches check
  # "original" patches (frequently used) can only be checked with a ram dumb of the game
  # the folder "ram_dumps" will hold wii game ram dumps (24 mb approx) named as the game ID
  # if a patch original condition fails it is simply not applied so it isn't an error to be checked
  
  # everything is good for now
  return True

# function to be able to remove the tmp folder
# after the check_riiv_patches() function check
def check_riiv_patches_wrapper(game_path, xml_path, mod_files_folder):
    
  # safe the return value of check_riiv_patches()
  result = check_riiv_patches(game_path, xml_path, mod_files_folder)
  file_ops.rm_folder("tmp")
  return result

# get_riiv_patches_inf() function
# function to return information about the patches used in the riivolution file
def get_riiv_patches_inf(game_path, xml_path, mod_files_folder):  
  
  # first check if the XML is well formatted
  # if the game relates to the XML
  # and if the mod files folder/xml file/game file relate to each other
  if (check_riiv_patches_wrapper(game_path, xml_path, mod_files_folder) == False):
    return None
  
  # to be able to store correctly the XML information
  # literally the first time doing classes stuff (*jesus what is this*)
  class section:
      def __init__(self, name):
        self.name = name
        self.options = []
  class option:
      def __init__(self, name):
        self.name = name
        self.choices = []
  class choice:
      def __init__(self, name):
        self.name = name
        self.patches = []
  class patch:
      def __init__(self, id):
        self.id = id
  
  # parse the XML and find the options element  
  # return a list of section classes with their respective option -> choice
  sections = []
  for sec in etree.parse(xml_path, etree.XMLParser(remove_comments = True)).find("options"):
    sections.append(section(sec.attrib["name"]))
    for opt in sec:
      sections[-1].options.append(option(opt.attrib["name"]))
      for cho in opt:
        sections[-1].options[-1].choices.append(choice(cho.attrib["name"]))
        for pat in cho:
          sections[-1].options[-1].choices[-1].patches.append(patch(pat.attrib["id"]))
  
  # ~ print()
  # ~ for sec in sections:
    # ~ print("Section name: %s" % (sec.name))
    # ~ print("#-Options names:")
    # ~ for opt in sec.options:
      # ~ print("----%s" % (opt.name))
      # ~ print("##----Choices names:")
      # ~ for cho in opt.choices:
        # ~ print("--------%s" % (cho.name))
        # ~ print("###-------Patches ids:")
        # ~ for pat in cho.patches:
          # ~ print("------------%s" % (pat.id))
  # ~ print()
  
  # return dis shit
  return sections

# apply_riiv_patches() function
def apply_riiv_patches(game_path, xml_path, mod_files_folder, choice_str_list):  
  
  # first check everything
  if (check_riiv_patches_wrapper(game_path, xml_path, mod_files_folder) == False):
    return False
  
  # check if choice_str_list contains valid XML choices
  # parse the XML and find the options element
  xml = etree.parse(xml_path, etree.XMLParser(remove_comments = True))
  root = xml.getroot()
  choices_names = []
  for sec in root.find("options"):
    for opt in sec:
      for cho in opt:
        choices_names.append(cho.attrib["name"])
  for choice_str in choice_str_list:
    if (choice_str not in choices_names):
      print("Invalid choice selected: %s" % (choice_str))
      return False
  
  # start the patching process
  
  # get the patches ids to apply
  patches_ids = []
  for sec in root.find("options"):
    for opt in sec:
      for cho in opt:
        if (cho.attrib["name"] in choice_str_list):
          for patch in cho:
            patches_ids.append(patch.attrib["id"])
  # get the patch elements on the xml
  patch_elems = []
  for patch in root.findall("patch"):
    if (patch.attrib["id"] in patches_ids):
      patch_elems.append(patch)
  
  # get game id code pieces
  game_id = ["", "", ""]
  game_id[0] = exec_subprocess([wit_path, "ID6", game_path]).stdout.replace("\n", "")
  game_id[1] = game_id[0][3:-2]
  game_id[2] = game_id[0][4:]
  game_id[0] = game_id[0][:-3]
  
  # extract the files from the game (wish I would know a way to just
  # have to deal only with the mod files and somehow "merge them" on
  # the original WBFS so no complete extraction has to be done)
  result = exec_subprocess([wit_path, "extract", game_path, "tmp/"]).returncode
  if (result != 0):
    print("wit could not extract correctly the game files")
    return False
  
  # define the actual root of files to change (its either "tmp/" or "tmp/DATA/")
  ext_game_path = "tmp/files/"
  if (file_ops.f_exists("tmp/DATA/")):
    ext_game_path = "tmp/DATA/files/"
  
  # First, deal with the file/folder replacements
  memory_original_exists = False # see if the original condition exists in a memory patch
  for patch in patch_elems:
    # get the root path of the current pacthes
    patch_root_path = rtn_def_param_str(file_ops.get_base_path(mod_files_folder, True) + "/" 
                                        + get_patch_elem_root_path(root, patch), game_id)
    # go through each patch type
    for patch_type in patch:       
      # file patch
      if (patch_type.tag == "file"):
        fpath1 = patch_root_path + "/" + rtn_def_param_str(patch_type.attrib["external"], game_id)
        fpath2 = rtn_def_param_str(ext_game_path + patch_type.attrib["disc"], game_id)
        file_ops.cp_file(fpath1, fpath2)
      # folder patch
      if (patch_type.tag == "folder"):
        # check if the folder in the extracted files exists
        if ("disc" in patch_type.attrib):
          fpath1 = patch_root_path + "/" + rtn_def_param_str(patch_type.attrib["external"], game_id)
          fpath2 = rtn_def_param_str(ext_game_path + patch_type.attrib["disc"], game_id)
          file_ops.cp_folder(fpath1, fpath2, False)
      # memory patch
      if (patch_type.tag == "memory"):
        if ("original" in patch_type.attrib):
          memory_original_exists = True
      
  
  # Now, handle the memory patches
  # wit dolpatch and geckoloader will be used
  # - wit will patch offsets above 0x8000FFFF
  # - geckoloader will patch offsets from 0x8000FFFF and below
  # for the wit patches a temporal riivolution xml will be made
  # for the geckoloader patches a temporal gecko code list will be done
  # all patches will be checked agaisnt a ram dump of the game
  ram_dump = None
  if (memory_original_exists): # open the ram dump
    ram_dump = open("ram_dumps/" + game_id[0] + game_id[1] + game_id[2] + ".bin", "rb")
  tmp_gl = open("tmp_gl.txt", "w")
  tmp_wit = open("tmp_wit.xml", "w")
  
  # start the loop!
  for patch in patch_elems:
    patch_root_path = rtn_def_param_str(file_ops.get_base_path(mod_files_folder, True) + "/" 
                                        + get_patch_elem_root_path(root, patch), game_id)
    for patch_type in patch:
      # only go through the memory patches this time
      if (patch_type.tag != "memory"):
        continue
      offset_hex_str = patch_type.attrib["offset"][2:]
      offset_hex_int = int(offset_hex_str, 16)
      
      # check the original value (if it has one)
      if ("original" in patch_type.attrib):
        og_hex = patch_type.attrib["original"]
        ram_dump.seek(offset_hex_int - 0x80000000)
        dmp_hex = ram_dump.read(int(len(og_hex) / 2)).hex()
        # ignore the patch if original does not match
        if (og_hex.upper() != dmp_hex.upper()):
          continue
        # wit hotfix (report this inconsistency)
        patch_type.attrib["original"] = patch_type.attrib["original"].replace("0x", "")
      
      # check if the memory patch contains a valuefile attribute
      if ("valuefile" in patch_type.attrib):
        val_file_path = patch_root_path + "/" + rtn_def_param_str(patch_type.attrib["valuefile"], game_id)
        tmp = open(val_file_path, "rb")
        del patch_type.attrib["valuefile"]
        patch_type.set("value", tmp.read(file_ops.get_file_size(val_file_path)).hex().upper())
        tmp.close()
      else: # it has a value attribute
        # wit hotfix (report this inconsistency)
        patch_type.attrib["value"] = patch_type.attrib["value"].replace("0x", "")
      
      # for wit
      if (offset_hex_int > 0x8000FFFF):
        tmp = etree.tostring(patch_type).decode("utf-8")
        while (tmp[0] in " \t\r\n"):
          tmp = tmp[1:]
        while (tmp[-1] in " \t\r\n"):
          tmp = tmp[:-1]
        tmp_wit.write(tmp + "\n")
      # for geckoloader
      else:
        tmp = patch_type.attrib["value"]
        last_offset = offset_hex_int # variable to be used later
        if (len(tmp) > 8):
          for i in range(0, int(len(tmp) / 8)):
            last_offset = offset_hex_int + (4 * i)
            if (len(tmp) > 8):
              tmp_gl.write("%08X %s\n" % (last_offset - 0x80000000 + 0x04000000, tmp[:-len(tmp) + 8]))
            else:
              tmp_gl.write("%08X %s\n" % (last_offset - 0x80000000 + 0x04000000, tmp))
            tmp = tmp[8:]
          last_offset += 4
        # write the rest of tmp depending on its size
        if (len(tmp) == 6):
          tmp_gl.write("%08X 0000%s\n" % (last_offset - 0x80000000 + 0x02000000, tmp[:-2]))
          tmp_gl.write("%08X 000000%s\n" % (last_offset + 2 - 0x80000000, tmp[4:]))
        elif (len(tmp) == 4):
          tmp_gl.write("%08X 0000%s\n" % (last_offset - 0x80000000 + 0x02000000, tmp))
        elif (len(tmp) == 2):
          tmp_gl.write("%08X 000000%s\n" % (last_offset - 0x80000000, tmp))
          
  # close files
  tmp_wit.close()
  tmp_gl.close()
  if (memory_original_exists):
    ram_dump.close()
  
  # define the main.dol location
  dol_path = ext_game_path.replace("files/", "sys/main.dol")
  dol_tmp_path = dol_path.replace("main.dol", "main_tmp.dol")
  
  # patch the forsaken main.dol
  # geckoloader part
  if (file_ops.get_file_size("tmp_gl.txt") != 0):
    # check the return code from geckoloader
    result = exec_subprocess(geckoloader_path + [dol_path, "tmp_gl.txt", "--dest",  
                             dol_tmp_path, "--txtcodes", "ALL", "--optimize"]).returncode
    if (result != 0):
      print("geckoloader could not handle tmp_gl.txt properly")
      return False
    file_ops.cp_file(dol_tmp_path, dol_path)
    file_ops.rm_file(dol_tmp_path)
  # wit part
  if (file_ops.get_file_size("tmp_wit.xml") != 0):
    # check the return code from wit
    result = exec_subprocess([wit_path, "dolpatch", dol_path,
                             "XML=tmp_wit.xml", "--dest", dol_tmp_path]).returncode
    if (result != 0):
      print("wit could not handle tmp_wit.xml properly")
      return False
    file_ops.cp_file(dol_tmp_path, dol_path)
    file_ops.rm_file(dol_tmp_path)
  
  # build game image file
  exec_subprocess([wit_path, "copy", "tmp/", "result.wbfs"])
  
  # eliminate temp files
  file_ops.rm_folder("tmp/")
  file_ops.rm_file("tmp_gl.txt")
  file_ops.rm_file("tmp_wit.xml")
  
  # done! (hopefully)
  return True

# apply_riiv_patches_wrapper() function
def apply_riiv_patches_wrapper(game_path, xml_path, mod_files_folder, choice_str_list):
  
  # safe the return value of apply_riiv_patches()
  result = apply_riiv_patches(game_path, xml_path, mod_files_folder, choice_str_list)
  file_ops.rm_folder("tmp/")
  file_ops.rm_file("tmp_gl.txt")
  file_ops.rm_file("tmp_wit.xml")
  return result
